// ProcessGuard.cpp : 此文件包含 "main" 函数。程序执行将在此处开始并结束。
//

#include <iostream>
#include <Windows/WinProcess.hpp>
#include <Windows/WinSystem.hpp>
#include <filesystem>
#include <iostream>
#include <string>
#include "ThreadPool.h"

// 命名空间别名简化
namespace WinProcess = CppWinApi::Windows::WinProcess;
namespace WinSystem = CppWinApi::Windows::WinSystem;
namespace fs = std::filesystem;

// 函数声明
std::string MakeUniqueName(const std::string& uniqueId, const std::string& name, bool isGlobal = false);

// [时序1] 启动函数 - 由初始进程调用
void Start(const std::string& workProcessExe);
// [时序2] 工作进程保护函数 - 由守护进程调用
void ProtectingWorkProcess(const std::string& workProcessExePath);
// [时序3] 子进程保护函数 - 由守护进程调用
void ProtectingChildProcess(const std::string& workProcessExe, bool isStartFromChildProcess);
// [时序4] 守护进程保护函数 - 由子进程调用
void ProctetingGuardProcess(const std::string& workProcessExe);

// 全局唯一标识符，用于创建命名的互斥体
const char PROCESS_GUARD_ID[] = "13604461032@163.COM_PROCESS_GUARD";

// 全局互斥体句柄
HANDLE gWorkProcessMutex = nullptr;     // 工作进程互斥体
HANDLE gProcessGuardMutex = nullptr;    // 守护进程互斥体
HANDLE gChildProcessMutex = nullptr;    // 子进程互斥体

// 全局线程池指针
AsyncThread::ThreadPool* gThreadPool;

int main(int argc, char** args)
{
	std::cout << "Process Id: " << ::GetCurrentProcessId() << "\n";

	// [时序0] 程序入口，解析命令行参数
	// 约定首个参数是命令
	auto command = std::string(args[1]);

	if (command == "start")
	{
		// [时序1] 启动模式
		std::cout << "工作进程：" << args[2] << "\n";
		Start(args[2]);  // 调用启动函数
	}
	else if (command == "launch-process-guard")
	{
		// [时序2] 守护进程模式
		gThreadPool = new AsyncThread::ThreadPool(2);  // 创建2个线程的线程池

		// 创建或打开命名的互斥体
		gWorkProcessMutex = WinSystem::OpenOrCreateMutex(MakeUniqueName(PROCESS_GUARD_ID, "work_process_mutex"));
		gProcessGuardMutex = WinSystem::OpenOrCreateMutex(MakeUniqueName(PROCESS_GUARD_ID, "process_guard_mutex"));
		gChildProcessMutex = WinSystem::OpenOrCreateMutex(MakeUniqueName(PROCESS_GUARD_ID, "process_guard_child_mutex"));

		// 获取守护进程的互斥锁，表示当前进程是活动的守护进程
		::WaitForSingleObject(gProcessGuardMutex, INFINITE);

		bool isStartFromChildProcess = false;

		if (argc == 4)
		{
			auto temp = std::string(args[3]);
			isStartFromChildProcess = temp == "is-start-from-child-process";

			if (!isStartFromChildProcess)
			{
				// 释放守护进程互斥锁，通知子进程可以退出
				::ReleaseMutex(gProcessGuardMutex);

				MessageBoxA(nullptr, "请使用正确的参数启动守护进程！\n"
					"正确格式：launch-process-guard <工作进程路径> [is-start-from-child-process]", "守护进程提示", MB_OK | MB_ICONERROR);
				return 1;
			}
		}

		// [时序3] 使用线程池异步执行保护子进程的任务
		auto task1 = gThreadPool->Enqueue([args, argc, isStartFromChildProcess]() {
			ProtectingChildProcess(args[2], isStartFromChildProcess);
			});

		// [时序2] 使用线程池异步执行保护工作进程的任务
		auto task2 = gThreadPool->Enqueue([args]() {
			ProtectingWorkProcess(args[2]);
			});

		// 等待工作进程保护任务完成（工作进程主动退出）
		task2.wait();

		// 释放守护进程互斥锁，通知子进程可以退出
		::ReleaseMutex(gProcessGuardMutex);

		// 等待子进程退出
		task1.wait();

		// 清理资源
		CloseHandle(gWorkProcessMutex);
		CloseHandle(gProcessGuardMutex);
		CloseHandle(gChildProcessMutex);
		delete gThreadPool;
	}
	else if (command == "launch-process-guard-child")
	{
		// [时序4] 子进程模式
		gProcessGuardMutex = WinSystem::OpenOrCreateMutex(MakeUniqueName(PROCESS_GUARD_ID, "process_guard_mutex"));
		gChildProcessMutex = WinSystem::OpenOrCreateMutex(MakeUniqueName(PROCESS_GUARD_ID, "process_guard_child_mutex"));

		// 调用守护进程保护函数
		ProctetingGuardProcess(args[2]);

		// 清理资源
		::CloseHandle(gProcessGuardMutex);
		::CloseHandle(gChildProcessMutex);
	}

	return 0;
}

void ProctetingGuardProcess(const std::string& workProcessExe)
{
	std::cout << "开始保护守护进程...\n";

	// [时序4.1] 获取子进程互斥锁
	::WaitForSingleObject(gChildProcessMutex, INFINITE);

	auto startPath = WinProcess::GetProcessExePath(::GetCurrentProcessId());
	// 准备启动守护进程的参数
	auto args = std::string("launch-process-guard ") + "\"" + workProcessExe + "\" " + "is-start-from-child-process";

	do {
		// [时序4.2] 监视守护进程状态
		auto result = WaitForSingleObject(gProcessGuardMutex, INFINITE);

		if (result == WAIT_FAILED)
		{
			// [时序4.3] 错误处理
			::ReleaseMutex(gChildProcessMutex);
			auto error = "等待守卫进程错误, 错误码：" + std::to_string(::GetLastError());
			::MessageBoxA(nullptr, error.c_str(), "守护进程提示", 0);
			break;
		}
		else if (result == WAIT_OBJECT_0)
		{
			// [时序4.4] 正常退出
			::ReleaseMutex(gChildProcessMutex);
			break;
		}

		// [时序4.5] 释放锁以便守护进程可以获取
		::ReleaseMutex(gProcessGuardMutex);

		std::cout << "启动守护进程...\n";

		// [时序4.6] 重启守护进程
		//auto processInfo = WinProcess::StartProcess(startPath.string(), CREATE_NEW_CONSOLE, args, startPath.parent_path().string());
		auto processInfo  = WinProcess::StartBackgroundProcess(startPath.string(), args, startPath.parent_path().string());
		auto dwMillsecond = 1000;

		do {
			// [时序4.7] 等待守护进程启动
			DWORD code = 0;
			auto status = processInfo.WaitForExit(&code, dwMillsecond);

			if (status == WAIT_OBJECT_0)
			{
				// [时序4.8] 守护进程异常退出
				::ReleaseMutex(gChildProcessMutex);
				return;
			}

			// [时序4.9] 检查守护进程是否获取了锁
			status = ::WaitForSingleObject(gProcessGuardMutex, 0);

			if (status == WAIT_TIMEOUT)
			{
				// [时序4.10] 守护进程已正常运行
				break;
			}

			// [时序4.11] 释放锁并增加等待时间
			::ReleaseMutex(gProcessGuardMutex);
			dwMillsecond *= 2;
		} while (true);

	} while (true);
}

void ProtectingChildProcess(const std::string& workProcessExe, bool isStartFromChildProcess)
{
	std::cout << "开始保护子进程...\n";

	auto startPath = WinProcess::GetProcessExePath(::GetCurrentProcessId());
	auto args = std::string("launch-process-guard-child ") + "\"" + workProcessExe + "\"";
	bool isFirstStart = true;

	do {
		// [时序3.1] 监视子进程状态
		auto result = ::WaitForSingleObject(gChildProcessMutex, INFINITE);

		// result == WAIT_OBJECT_0 的情况有两种：
		// 1、如果是从保护进程启动
		//	* 首次 WaitForSingleObject 一定是 WAIT_OBJECT_0，因为子进程还未启动
		//  * 此时 WAIT_OBJECT_0 不是由子进程向主进程发出退出程序信号
		//  * 我们不应该退出程序，而是立即去创建子进程
		// 2、如果是从子进程启动
		//  * 此时子进程已经启动，子进程主动退出时会发出退出信号，此时 WAIT_OBJECT_0 由子进程发出
		//  * 此时我们应该退出程序，避免重复创建子进程

		if (result == WAIT_OBJECT_0 && (!isFirstStart || isStartFromChildProcess))
		{
			// [时序3.2] 子进程主动退出
			break;
		}

		if (isFirstStart)
		{
			isFirstStart = false;
		}

		if (result == WAIT_FAILED)
		{
			// [时序3.3] 错误处理
			auto error = "等待子进程错误, 错误码：" + std::to_string(::GetLastError());
			::MessageBoxA(nullptr, error.c_str(), "守护进程提示", 0);
			break;
		}

		// [时序3.4] 释放锁以便子进程可以获取
		::ReleaseMutex(gChildProcessMutex);

		std::cout << "启动子进程...\n";

		// [时序3.5] 启动子进程
		//auto processInfo = WinProcess::StartProcess(startPath.string(), CREATE_NEW_CONSOLE, args, startPath.parent_path().string());
		auto processInfo = WinProcess::StartBackgroundProcess(startPath.string(), args, startPath.parent_path().string());
		auto dwMillsecond = 1000;

		do {
			// [时序3.6] 等待子进程启动
			DWORD code = 0;
			auto status = processInfo.WaitForExit(&code, dwMillsecond);

			if (status == WAIT_OBJECT_0)
			{
				return;
			}

			// [时序3.7] 检查子进程是否获取了锁
			status = ::WaitForSingleObject(gChildProcessMutex, 0);
			if (status == WAIT_TIMEOUT)
			{
				// [时序3.8] 子进程已正常运行
				break;
			}

			// [时序3.9] 释放锁并增加等待时间
			::ReleaseMutex(gChildProcessMutex);
			dwMillsecond *= 2;
		} while (true);

	} while (true);
}

void ProtectingWorkProcess(const std::string& workProcessExePath)
{
	std::cout << "开始保护工作进程...\n";

	fs::path exePath = workProcessExePath;

	do {
		// [时序2.1] 监视工作进程状态
		auto result = ::WaitForSingleObject(gWorkProcessMutex, INFINITE);

		if (result == WAIT_OBJECT_0)
		{
			// [时序2.2] 工作进程主动退出
			break;
		}
		else if (result == WAIT_FAILED)
		{
			// [时序2.3] 错误处理
			auto error = "等待工作进程错误, 错误码：" + std::to_string(::GetLastError());
			::MessageBoxA(nullptr, error.c_str(), "守护进程提示", 0);
			break;
		}

		// [时序2.4] 释放锁以便工作进程可以获取
		::ReleaseMutex(gWorkProcessMutex);

		std::cout << "启动工作进程...\n";

		// [时序2.5] 启动工作进程
		auto processInfo = WinProcess::StartProcess(exePath.string(), 0, "start-from-process-guard", exePath.parent_path().string());
		auto dwMillsecond = 1000;

		do {
			// [时序2.6] 等待工作进程启动
			DWORD code = 0;
			auto status = processInfo.WaitForExit(&code, dwMillsecond);

			if (status == WAIT_OBJECT_0)
			{
				// [时序2.7] 工作进程异常退出
				auto error = "工作进程闪退, 退出码：" + std::to_string(code);
				::MessageBoxA(nullptr, error.c_str(), "守护进程提示", 0);
				return;
			}

			// [时序2.8] 检查工作进程是否获取了锁
			status = ::WaitForSingleObject(gWorkProcessMutex, 0);
			if (status == WAIT_TIMEOUT)
			{
				// [时序2.9] 工作进程已正常运行
				break;
			}

			// [时序2.10] 释放锁并增加等待时间
			::ReleaseMutex(gWorkProcessMutex);
			dwMillsecond *= 2;
		} while (true);
	} while (true);
}

void Start(const std::string& workProcessExe)
{
	// [时序1.1] 获取当前进程路径
	auto startPath = WinProcess::GetProcessExePath(::GetCurrentProcessId());

	// [时序1.2] 准备启动参数
	auto args = std::string("launch-process-guard ") + "\"" + workProcessExe + "\"";

	// [时序1.3] 启动守护进程
	//WinProcess::StartProcess(startPath.string(), CREATE_NEW_CONSOLE, args, startPath.parent_path().string());
	WinProcess::StartBackgroundProcess(startPath.string(), args, startPath.parent_path().string());
}

std::string MakeUniqueName(const std::string& uniqueId, const std::string& name, bool isGlobal)
{
	// 创建全局或本地命名的互斥体名称
	if (isGlobal)
	{
		return std::string("Global\\") + uniqueId + "_" + name;
	}
	else
	{
		return std::string("Local\\") + uniqueId + "_" + name;
	}
}